# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import functools
import logging
import os
import time

from absl import app, flags

from coin.base.book.util import convert_book_to_dict, convert_trade_to_dict
from coin.exchange.base.rest_order_gateway_base import pretty_json
from coin.experimental.xguo.fruit.base.event_clock import event_clock
from coin.experimental.xguo.fruit.base.strategy_config import BaseStrategyConfig
from coin.experimental.xguo.fruit.util.record_writer import RecordWriter
from coin.experimental.xguo.fruit.util.theory import run_from_generated_feed
from coin.experimental.xguo.fruit.util.util import guess_tick_info
from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.subscription import get_subscription_key
from coin.experimental.xguo.fruit.base.run_from_archive import (
    run_from_archive_or_cache, define_fastfeed_flags)

from coin.exchange.bithumb.kr_rest.product_info import bithumb_tick_info
from coin.exchange.coinone_v2.kr_rest.product_info import coinone_tick_info
from coin.exchange.upbit_v1.kr_rest.product_info import upbit_tick_info

FLAGS = flags.FLAGS


def is_same_me(subreq1, subreq2):
  x = subreq1.market_type == subreq2.market_type
  y = subreq1.exchange == subreq2.exchange
  return x and y


class BaseSimStrategy(object):
  def __init__(self, strategy_config: BaseStrategyConfig):
    self._logger = logging.getLogger(__name__)
    self._strategy_config = strategy_config
    self._executors = {}
    self._writer = RecordWriter(filename='smm_sim_record_%s.gz' % int(time.time() * 1e9))
    self._guess_tick_info_done = False

  @property
  def executors(self):
    return self._executors

  @property
  def strategy_config(self):
    return self._strategy_config

  @property
  def product_list(self):
    config_dict = self._strategy_config.feed_subsystem_config
    products = []
    if config_dict is None:
      self._logger.info('Feed system config not found!')
      return products

    for config in config_dict.values():
      products.extend(config.products)
    return products

  @property
  def currency_list(self):
    currency_set = set()
    for product in self.product_list:
      currency_set.add(product.base)
      currency_set.add(product.quote)
    return list(currency_set)

  def on_book_reset(self, products, book_builder_name, book_builder, feed_cache_key=None):
    for product in products:
      subreq = get_subscription_key(product)
      if is_same_me(book_builder_name, subreq):
        book_builder.subscribe_book(product, functools.partial(self.common_on_book, product))
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def prepare_feed_subsystem(self):
    config_dict = self._strategy_config.feed_subsystem_config
    if config_dict is None:
      self._logger.error('Feed subsystem config is not found!')
      return

    sim_config = self._strategy_config.sim_config
    if sim_config is None:
      self._logger.error('Simulation config is not found!')
      return

    if sim_config.get('use_gen_feed'):
      run_from_generated_feed(self.product_list, self)
    elif sim_config.get('use_feed_cache', True):
      define_fastfeed_flags()
      FLAGS.feed_machine = sim_config['machine']
      run_from_archive_or_cache(sim_config['from_ts'], sim_config['machine'], self.product_list, self, True)
    else:
      run_from_archive(self.product_list,
                       functools.partial(self.on_book_reset, self.product_list),
                       from_ts=sim_config['from_ts'],
                       to_ts=sim_config['to_ts'],
                       machine=sim_config['machine'],
                       root_dir=sim_config.get('root_dir'),
                       endpoint_url=None,
                       worker_id='1',
                       feed_checker_config=None,
                       fs_config=None,
                       use_feed_arbitration=None,
                       skip_failed_product=False)

  def log_json_message(self, data, *, msg='', level='INFO'):
    assert isinstance(data, (list, dict)), type(data)
    if isinstance(level, str):
      level = logging.getLevelName(level)
    self._logger.log(level, '%s:\n%s\n', msg, pretty_json(data))

  def log_sim_time(self):
    if event_clock.is_ready():
      self._logger.info('Current time: %s', event_clock.now())
    else:
      self._logger.info('No feed received yet!')

  def common_on_book(self, product, book):
    event_clock.update(book.timestamp)
    if not self._guess_tick_info_done:
      if product.exchange == 'Bithumb':
        bithumb_tick_info[product.symbol] = guess_tick_info(book)
      elif product.exchange == 'Coinone':
        coinone_tick_info[product.symbol] = guess_tick_info(book)
      elif product.exchange == 'Upbit':
        upbit_tick_info[product.symbol] = guess_tick_info(book)
      else:
        pass

      num_bithumb = 0
      num_coinone = 0
      num_upbit = 0
      for prod in self.product_list:
        if prod.exchange == 'Bithumb':
          num_bithumb += 1
        elif prod.exchange == 'Coinone':
          num_coinone += 1
        elif prod.exchange == 'Upbit':
          num_upbit += 1
      a = num_bithumb == len(bithumb_tick_info)
      b = num_coinone == len(coinone_tick_info)
      c = num_upbit == len(upbit_tick_info)
      if a and b and c:
        self._guess_tick_info_done = True

    self.on_book(product, book)

  def on_book(self, product, book):
    book_dict = convert_book_to_dict(book, level=10, notional_qty=10)
    self.log_json_message(book_dict, msg='book')

  def on_trade(self, product, trade_event):
    trade_dict = convert_trade_to_dict(product, trade_event)
    self.log_json_message(trade_dict, msg='trade')

  def on_fill(self, price_last_fill, qty_last_fill, order):
    # noinspection PyProtectedMember
    fill_dict = {
        'price_last_fill': price_last_fill,
        'qty_last_fill': qty_last_fill,
        'order': order._asdict()
    }
    self.log_json_message(fill_dict, msg='fill')

  def run(self):
    self.prepare_feed_subsystem()


def main(_):
  if FLAGS.config is None:
    config_file = os.path.join(os.path.dirname(__file__), 'sim_config.yaml')
  else:
    config_file = FLAGS.config
  config = BaseStrategyConfig.from_file(config_file)
  strategy = BaseSimStrategy(strategy_config=config)
  strategy.run()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
